import numpy as np
import networkx as nx
from sklearn.ensemble import IsolationForest
#时序异构图异常检测服务 (THGAD)
class TemporalHeterogeneousGraphAnomalyDetection:
    def __init__(self, window_size=5, contamination=0.1):
        self.window_size = window_size
        self.contamination = contamination
        self.models = {}  # 为每种节点类型存储一个模型
    
    def extract_features(self, graph_sequence, node_type):
        """从图序列中提取时序特征"""
        features = []
        for node in graph_sequence[-1].nodes():
            if graph_sequence[-1].nodes[node]['type'] == node_type:
                # 提取度变化特征
                degree_features = []
                for t in range(max(0, len(graph_sequence)-self.window_size), len(graph_sequence)):
                    if node in graph_sequence[t]:
                        degree_features.append(graph_sequence[t].degree(node))
                    else:
                        degree_features.append(0)
                
                # 提取邻居类型分布特征
                type_dist = {}
                for neighbor in graph_sequence[-1].neighbors(node):
                    n_type = graph_sequence[-1].nodes[neighbor]['type']
                    type_dist[n_type] = type_dist.get(n_type, 0) + 1
                
                # 合并特征
                node_features = degree_features + list(type_dist.values())
                features.append((node, node_features))
        
        return features
    
    def fit(self, graph_sequence):
        """训练异常检测模型"""
        node_types = set()
        for node in graph_sequence[-1].nodes():
            node_types.add(graph_sequence[-1].nodes[node]['type'])
        
        for node_type in node_types:
            features = self.extract_features(graph_sequence, node_type)
            if not features:
                continue
                
            nodes, X = zip(*features)
            model = IsolationForest(contamination=self.contamination)
            model.fit(X)
            self.models[node_type] = {'model': model, 'nodes': nodes}
    
    def detect(self, graph_sequence):
        """检测异常"""
        anomalies = {}
        for node_type, model_info in self.models.items():
            features = self.extract_features(graph_sequence, node_type)
            if not features:
                continue
                
            nodes, X = zip(*features)
            scores = model_info['model'].decision_function(X)
            for i, node in enumerate(nodes):
                if scores[i] < 0:  # 负分数表示异常
                    anomalies[node] = scores[i]
        
        return anomalies 